#!/usr/bin/env python3
# ui/main_window.py
# 主窗口：布局、按钮、更新逻辑
import sys
import time
import numpy as np
from PyQt5.QtWidgets import (
    QApplication,
    QMainWindow,
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QSplitter,
    QPushButton,
    QLabel,
)
from PyQt5.QtCore import QTimer, Qt
from PyQt5.QtGui import QFont

from .video_server import VideoServer, VideoPanel
from .map_widget import GraphMapWidget
from .info_table import RobotInfoTable
from map_manager.obstacle_editor import ObstacleEditor  # 导入编辑器
from map_manager.map_manage import MapManager

import logging

logger = logging.getLogger(__name__)


class PursuitGameUI(QMainWindow):
    def __init__(self, server=None, logger=None):
        super().__init__()
        if logger is None:
            logger = logging.getLogger(__name__)
        self.logger = logger
        self.server = server
        self.setWindowTitle("追捕任务监控系统")
        self.setGeometry(100, 100, 1200, 800)
        # 视频服务器
        self.video_server = VideoServer()
        self.video_panels = [
            VideoPanel(self.video_server, 1),
            VideoPanel(self.video_server, 2),
            VideoPanel(self.video_server, 3),
        ]
        self.prev_clients = set()
        # 主布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        main_splitter = QSplitter(Qt.Horizontal)
        # 左侧：地图和表格
        left_widget = QWidget()
        left_layout = QVBoxLayout(left_widget)
        if server and hasattr(server, "map_manager"):
            self.map_manager = server.map_manager
            self.map_widget = GraphMapWidget(server.map_manager)
        else:
            self.map_manager = MapManager(logger)
            self.map_widget = GraphMapWidget(self.map_manager)
        left_layout.addWidget(self.map_widget, 1)
        left_layout.addStretch()
        self.robot_table = RobotInfoTable()
        self.robot_table.setFixedHeight(150)
        left_layout.addWidget(self.robot_table)
        # 右侧：视频
        right_widget = QWidget()
        right_layout = QVBoxLayout(right_widget)
        for panel in self.video_panels:
            right_layout.addWidget(panel)
        main_splitter.addWidget(left_widget)
        main_splitter.addWidget(right_widget)
        main_splitter.setSizes([800, 400])
        main_layout.addWidget(main_splitter)
        # 控制布局
        control_status_layout = QHBoxLayout()
        control_status_layout.setContentsMargins(0, 5, 0, 5)
        self.pause_btn = QPushButton("暂停任务")
        self.pause_btn.clicked.connect(self.toggle_pause)
        control_status_layout.addWidget(self.pause_btn)
        self.reset_btn = QPushButton("位置重置")
        self.reset_btn.clicked.connect(self.position_reset)
        control_status_layout.addWidget(self.reset_btn)
        self.edit_map_btn = QPushButton("地图编辑")
        self.edit_map_btn.clicked.connect(self.edit_map)
        control_status_layout.addWidget(self.edit_map_btn)
        self.reload_map_btn = QPushButton("重载地图")
        self.reload_map_btn.clicked.connect(self.reload_map)
        control_status_layout.addWidget(self.reload_map_btn)
        # 状态标签
        self.game_status_label = QLabel("任务状态: 等待连接")
        self.game_status_label.setFixedWidth(120)
        control_status_layout.addWidget(self.game_status_label)
        self.capture_label = QLabel("捕获次数: 0")
        self.capture_label.setFixedWidth(100)
        control_status_layout.addWidget(self.capture_label)
        self.belief_label = QLabel("信念: 支持位置 0 | 熵 0.00")
        self.belief_label.setFixedWidth(180)
        control_status_layout.addWidget(self.belief_label)
        self.online_label = QLabel("在线: 0/3")
        self.online_label.setFixedWidth(80)
        control_status_layout.addWidget(self.online_label)
        self.last_update_label = QLabel(f"更新: {time.strftime('%H:%M:%S')}")
        self.last_update_label.setFixedWidth(120)
        control_status_layout.addWidget(self.last_update_label)
        control_status_layout.addStretch()
        control_widget = QWidget()
        control_widget.setLayout(control_status_layout)
        control_widget.setFixedHeight(40)
        main_layout.addWidget(control_widget)
        # 状态栏
        self.status_bar = self.statusBar()
        self.status_bar.setFixedHeight(25)
        self.time_label = QLabel(time.strftime("%H:%M:%S"))
        self.status_bar.addPermanentWidget(self.time_label)
        # 更新定时器
        self.update_timer = QTimer(self)
        self.update_timer.timeout.connect(self.update_display)
        self.update_timer.start(200)

    def closeEvent(self, event):
        self.video_server.stop()
        super().closeEvent(event)

    def edit_map(self):
        """地图编辑"""
        logger.info(
            f"打开地图编辑器: {hasattr(self, 'map_manager')}   {self.map_manager}"
        )
        if hasattr(self, "map_manager") and self.map_manager:
            self.logger.info("打开地图编辑器")
            if self.map_manager.show_editor():
                if self.server:
                    self.server.policy_reload()
                self.reset_ui()

    def reset_ui(self):
        """重置UI"""
        if hasattr(self, "map_widget"):
            self.map_widget.reset()
        self.update_display()

    def toggle_pause(self):
        """切换暂停"""
        if self.server:
            self.server.toggle_game_active()

    def position_reset(self):
        """位置重置"""
        if self.server:
            self.server.position_reset()
            logger.info("位置重置已触发")

    def reload_map(self):
        """重新加载地图"""
        if self.server:
            self.server.policy_reload()
            logger.info("重载地图已触发")
            self.reset_ui()
        else:
            logger.warning("服务器未初始化")

    def update_display(self):
        """更新显示"""
        try:
            robots = []
            game_active = False
            capture_count = 0
            belief_pr = None
            belief_pos = None
            if self.server:
                if hasattr(self.server, "get_all_robots_info"):
                    robots = self.server.get_all_robots_info()
                if hasattr(self.server, "get_game_active"):
                    game_active = self.server.get_game_active()
                if hasattr(self.server, "get_capture_count"):
                    capture_count = self.server.get_capture_count()
                if hasattr(self.server, "get_belief_pr") and hasattr(
                    self.server, "get_belief_pos"
                ):
                    belief_pr = self.server.get_belief_pr()
                    belief_pos = self.server.get_belief_pos()
            self.pause_btn.setText("恢复任务" if game_active else "暂停任务")
            self.map_widget.update_robots(robots, capture_count, belief_pr, belief_pos)
            self.robot_table.update_robots(robots)
            game_status_text = (
                "<font color='green'>进行中</font>"
                if game_active
                else "<font color='orange'>暂停</font>"
            )
            self.game_status_label.setText(f"任务状态: {game_status_text}")
            self.capture_label.setText(f"捕获次数: {capture_count}")
            if belief_pos is not None and belief_pr is not None:
                support_count = np.sum(belief_pos)
                entropy = (
                    -np.sum(
                        belief_pr[belief_pos] * np.log2(belief_pr[belief_pos] + 1e-10)
                    )
                    if support_count > 0
                    else 0.0
                )
                self.belief_label.setText(
                    f"信念: 支持位置 {support_count} | 熵 {entropy:.2f}"
                )
            else:
                self.belief_label.setText("信念: 未初始化")
            online_count = sum(1 for r in robots if r.get("status") == "active")
            self.online_label.setText(f"在线: {online_count}/3")
            self.last_update_label.setText(f"更新: {time.strftime('%H:%M:%S')}")
            self.time_label.setText(time.strftime("%H:%M:%S"))
            # 自动分配视频源
            current_clients = set(self.video_server.clients.keys())
            new_addrs = current_clients - self.prev_clients
            self.prev_clients = current_clients
            for new_addr in sorted(new_addrs):
                for panel in self.video_panels:
                    if panel.selected_addr is None:
                        panel.set_source(new_addr)
                        break
        except Exception as e:
            logger.error(f"更新显示错误: {e}")
